package com.xuz.rbtdemo;

/**
 * 红黑树
 * 添加数据
 * 删除数据
 * 遍历查询数据
 * 要么为空，妖媚就有个Root节点
 *
 * @user 65453
 * @date 2024/9/2 22:43
 */
public class RBTree<K extends Comparable<K>, V> {

    // 定义红黑常量
    private static final boolean RED = false;
    private static final boolean BLACK = true;

    // 当前红黑树的Root节点
    private RBNode root;

    public RBNode getRoot() {
        return root;
    }

    public void setRoot(RBNode root) {
        this.root = root;
    }

    // 1.旋转和变色   2.添加数据【a.普通的二叉树插入 b.红黑的平衡】 3.删除数据  4.遍历数据

    /**
     * 左旋操作
     * 围绕p节点完成左旋操作
     * p               pr[r]
     * /\                /\
     * pl pr[r]  -->     p rr
     * /\            /\
     * rl rr         pl rl
     * 左旋的时候：
     * p-pl 和pr[r]-rr 这两根线时不变的
     * pr[r] - rl 需要改变为 p-rl
     */
    public void leftRotate(RBNode p) {
        if (p == null) {
            return;
        }
        // 获取旋转节点的右节点
        RBNode r = p.getRight();
        p.right = r.left;
        if (r.left != null) {
            // 表示如果r节点有左侧子节点，就需要挂载到旋转节点的右子节点
            r.left.parent = p;
        }
        // 不管p节点是否有父节点，都把这个父节点赋值给r的父节点
        r.parent = p.parent;
        if (p.parent == null) {
            this.root = r;
        } else if (p.parent.left == p) {
            // 说明p节点时其父节点的左侧子节点
            p.parent.left = r;
        } else {
            // 说明p节点是其父节点的右侧子节点
            p.parent.right = r;
        }
        // 设置p和r的关系
        r.left = p;
        p.parent = r;

    }

    /**
     * 右旋操作
     * 围绕p节点完成右旋操作
     * p               pl[l]
     * /\                /\
     * pl[l] pr  -->        ll p
     * /\                   /\
     * ll lr               lr pr
     * 右旋的时候：
     * p-pr 和pl[l]-ll 这两根线时不变的
     * pl[l] - lr 需要改变为 p-lr
     */
    public void rightRotate(RBNode p) {
        if (p == null) {
            return;
        }
        // 获取旋转节点的右节点
        RBNode l = p.getLeft();
        p.left = l.right;
        if (l.right != null) {
            // 表示如果l节点有右侧子节点，就需要挂载到旋转节点的左子节点
            l.right.parent = p;
        }
        // 不管p节点是否有父节点，都把这个父节点赋值给l的父节点
        l.parent = p.parent;
        if (p.parent == null) {
            this.root = l;
        } else if (p.parent.left == p) {
            // 说明p节点时其父节点的左侧子节点
            p.parent.left = l;
        } else {
            // 说明p节点是其父节点的右侧子节点
            p.parent.right = l;
        }
        // 设置p和r的关系
        l.left = p;
        p.parent = l;
    }

    /**
     * 新增节点的操作
     *
     * @param key
     * @param value
     */
    public void put(K key, V value) {
        if (key == null) {
            throw new NullPointerException();
        }
        // 1.向普通的二叉查找树中插入节点
        // 记录当前节点
        RBNode t = this.root;
        if (t == null) {
            // 说明整个红黑树是空的，那么插入的节点就是根节点(root)，root节点的parent为空
            this.root = new RBNode<>(key, value == null ? key : value, null);
            this.root.color = BLACK;
            return;
        }
        // 声明需要插入的节点的父节点
        RBNode parent;
        // 记录插入的位置 i<0 left;i>o right
        int i;
        // 右root节点，添加节点需要比较大小
        do {
            parent = t;
            // 输入的节点和当前比较的节点比较大小
            i = key.compareTo((K) t.key);
            if (i < 0) {
                t = t.left;
            } else if (i > 0) {
                t = t.right;
            } else {
                // 插入的节点与当前比较的节点的值相同
                // 覆盖掉原来的值
                t.setValue(value == null ? key : value);
                return;
            }
        } while (t != null);
        // 循环结束表示找到了需要插入节点的位置
        // 创建需要插入的节点
        RBNode<K, Object> e = new RBNode<>(key, value, parent);
        // 判断插入的节点实在parent的左侧还是右侧
        if (i < 0) {
            parent.left = e;
        } else {
            parent.right = e;
        }

        // 2. 调整平衡  旋转和变色
        fixAfterPut(e);

    }

    /**
     * 删除节点的操作
     *  1.节点的删除【普通的二叉查找树的节点删除】
     *  2.通过旋转和变色来实现平衡
     */
    public V remove(K key) {
        // 1.根据Key找到对应的RBNode在红黑树中的位置
        RBNode node = getNode(key);
        if (node == null) {
            return null;
        }
        V value = (V)node.getValue();
        deleteNode(node);
        return value;

    }

    /**
     * 根据Key找到对应的节点
     * @param key
     * @return
     */
    private RBNode getNode(K key){
        RBNode node = this.root;
        while (node != null) {
            int tmp = node.getKey().compareTo(key);
            if (tmp == 0) {
                return node;
            } else if (tmp > 0) {
                node = leftOf(node);
            } else {
                node = rightOf(node);
            }
        }
        return null;
    }

    /**
     * 具体删除节点的方法
     *  1.删除二叉树中的一个节点
     *      1).删除的是叶子节点  直接删除
     *      2).删除的节点有一个子节点  子节点替代删除的节点
     *      3).删除的节点有两个子节点  用前驱或者后继节点来替代
     *  2.删除节点后需要的旋转变色
     * @param node
     */
    private void deleteNode(RBNode node) {
        // 1).删除的节点有两个子节点  用前驱或者后继节点来替代
        if (leftOf(node) != null && rightOf(node) != null){
            // 找到删除节点的前驱或者后继节点
            RBNode predecessor= this.predecessor(node);
            // 找到前驱节点后，将前驱节点覆盖掉需要删除的节点
            node.key = predecessor.key;
            node.value = predecessor.value;
            // 需要删除的节点就从原来的node变成了前驱节点predecessor
            node = predecessor;

        }
        // 找到需要删除的节点
        RBNode replaceNode = leftOf(node) != null ? leftOf(node) : rightOf(node);
        if (replaceNode != null) {
            // 2).删除的节点有一个子节点  替代节点需要指向删除的节点
            replaceNode.parent = node.parent;
            if (parentOf(node) == null) {
                root = replaceNode;
            } else if (node == parentOf(node).left) {
                parentOf(node).left = replaceNode;
            } else {
                parentOf(node).right = replaceNode;
            }
            // 需要删除的节点起吊相关的关联关系  等待GC
            node.left = node.right = node.parent = null;
            // 考虑是否需要调整平衡
            if (colorOf(node) == BLACK) {
                // 需要做平衡的调整
                fixAfterRemove(replaceNode);
            }
        } else if (node.parent == null) {
            // 说明我们要删除的是root节点
            root = null;
        } else {
            // 2.
            if (colorOf(node) == BLACK) {
                // 需要平衡的调整
                fixAfterRemove(node);
            }
            if (parentOf(node).left == node) {
                parentOf(node).left = null;
            } else {
                parentOf(node).right = null;
            }
            node.parent = null;
            node = null;
        }

//        if (replaceNode != null) {
//            // 2).删除的节点有一个子节点  子节点替代删除的节点
//            node.key = replaceNode.key;
//            node.value = replaceNode.value;
//            // 需要删除的节点就从原来的node变成了前驱节点predecessor
//            node = replaceNode;
//
//        }
//        // 3).删除的是叶子节点  直接删除
//        RBNode p = parentOf(node);
//        if (leftOf(p) == node) {
//            p.left = null;
//        } else {
//            p.right = null;
//        }
//        node.parent = null;

    }

    /**
     * 删除节点做的平衡操作
     * 1.情况一：自己能搞定的，对应叶子节点是3节点和4节点
     * 2.情况二：自己搞不定，需要兄弟借，父亲节点下移，然后兄弟节点上移为新的父节点  变色＋旋转操作
     * 3.情况三：跟兄弟借，兄弟也没有  兄弟节点变红  递归处理
     * @param x
     */
    private void fixAfterRemove(RBNode x) {
        //处理情况2、3 删除节点是2节点
        while (x != root && colorOf(x) == BLACK) {
            // 删除节点是父亲节点的左侧子节点
            if (x == leftOf(parentOf(x))) {
                // 获取到当前节点的兄弟节点
                RBNode rbNode = rightOf(parentOf(x));
                if (colorOf(rbNode) == RED) { // 说明不是真正的兄弟节点  因为234数有左旋和右旋两种情况，所以可能不是真正的兄弟节点
                    // 变色＋左旋
                    setColor(rbNode, BLACK);
                    setColor(parentOf(x),RED);
                    leftRotate(parentOf(x));
                    // 找到正在的兄弟节点
                    rbNode = rightOf(parentOf(x));
                }
                // 情况三 兄弟节点没有任何子节点
                if (colorOf(leftOf(rbNode)) == BLACK && colorOf(rightOf(rbNode)) == BLACK) {
                    // 变色
                    setColor(rbNode,RED);
                    x = parentOf(x);
                } else {
                    // 情况二 兄弟节点借  旋转 变色
                    if (colorOf(rightOf(rbNode)) == BLACK) {
                        // 兄弟节点没有右节点，那么左节点肯定不为空
                        setColor(rbNode,RED);
                        setColor(leftOf(rbNode),BLACK);
                        rightRotate(rbNode);
                        // 旋转后，兄弟节点有了变化
                        rbNode = rightOf(parentOf(x));
                    }
                    // 左旋加变色的操作
                    setColor(rbNode, colorOf(parentOf(x)));
                    setColor(parentOf(x), BLACK);
                    setColor(rightOf(rbNode), BLACK);
                    // 做一次左旋操作
                    leftRotate(parentOf(x));
                    // 结束循环
                    x = root;
                }
            } else {
                // 删除节点是父亲节点的右侧子节点
                // 获取到当前节点的兄弟节点
                RBNode lbNode = leftOf(parentOf(x));
                if (colorOf(lbNode) == RED) { // 说明不是真正的兄弟节点  因为234数有左旋和右旋两种情况，所以可能不是真正的兄弟节点
                    // 变色＋右旋
                    setColor(lbNode, BLACK);
                    setColor(parentOf(x),RED);
                    rightRotate(parentOf(x));
                    // 找到正在的兄弟节点
                    lbNode = leftOf(parentOf(x));
                }
                // 情况三 兄弟节点没有任何子节点
                if (colorOf(leftOf(lbNode)) == BLACK && colorOf(rightOf(lbNode)) == BLACK) {
                    // 变色
                    setColor(lbNode,RED);
                    x = parentOf(x);
                } else {
                    // 情况二 兄弟节点借  旋转 变色
                    if (colorOf(leftOf(lbNode)) == BLACK) {
                        // 兄弟节点没有左节点，那么右节点肯定不为空
                        setColor(lbNode,RED);
                        setColor(rightOf(lbNode),BLACK);
                        leftRotate(lbNode);
                        // 旋转后，兄弟节点有了变化
                        lbNode = leftOf(parentOf(x));
                    }
                    // 左旋加变色的操作
                    setColor(lbNode, colorOf(parentOf(x)));
                    setColor(parentOf(x), BLACK);
                    setColor(leftOf(lbNode), BLACK);
                    // 做一次左旋操作
                    rightRotate(parentOf(x));
                    // 结束循环
                    x = root;
                }
            }
        }

        //处理情况一
        setColor(x, BLACK);

    }

    /**
     * 查询当前节点的前驱节点
     *  先找当前节点的左节点left，然后一直找左节点left的右节点
     * @param node
     * @return
     */
    private RBNode predecessor(RBNode node) {
        if (node == null) {
            return null;
        } else if (leftOf(node) != null) {
            RBNode leftNode = leftOf(node);
            while (rightOf(leftNode) != null) {
                leftNode = rightOf(leftNode);
            }
            return leftNode;
        } else {
            // 当前节点没有左侧子节点
            // 记录当前节点的父节点
            RBNode p = parentOf(node);
            RBNode ch = node;
            while (p != null && ch == leftOf(p)) {
                ch = p;
                p = p.parent;
            }
            return p;
        }
    }

    /**
     * 查询当前节点的后继节点
     *  先找当前节点的右节点right，然后一直找右节点right的左节点
     * @param node
     * @return
     */
    private RBNode successor(RBNode node) {
        if (node == null) {
            return null;
        } else if (rightOf(node) != null) {
            RBNode rightNode = rightOf(node);
            while (leftOf(rightNode) != null) {
                rightNode = leftOf(rightNode);
            }
            return rightNode;
        } else {
            // 当前节点没有右侧子节点
            // 记录当前节点的父节点
            RBNode p = parentOf(node);
            RBNode ch = node;
            while (p != null && ch == rightOf(p)) {
                ch = p;
                p = p.parent;
            }
            return p;
        }
    }

    /**
     * 插入节点后的调整处理
     * (可参考TreeMap中put方法的源码)
     * 1. 2-3-4树 新增元素 2节点添加一个元素将变为3节点 直接合并，节点中有两个元素
     *      红黑树：新增一个红色节点，这个红色节点会添加在黑色节点下(2节点) --- 这种情况不需要调整
     * 2. 2-3-4树 新增元素 3节点添加一个元素变为4节点合并 节点中有3个元素
     *      这里有6中情况，( 根左左 根左右  根右右 根右左)这四种要调整  (左中右的两种)不需要调整
     *      红黑树：新增红色节点 会添加到 上黑下红的节点中 = 排序后中间节点是黑色，两边节点是红色
     *
     *  3. 2-3-4树：新增一个元素 4节点添加一个元素需要裂变：中间元素升级为父节点，新增元素与剩下的其中一个合并
     *      红黑树：新增节点是红色+爷爷节点是黑色，父亲节点和叔叔节点为红色 调整为
     *              爷爷节点变红色，父亲和叔叔节点变为黑色，如果爷爷节点为root节点则调整为黑色
     * @param e
     */
    private void fixAfterPut(RBNode<K, Object> e) {
        if (e == null) {
            throw new NullPointerException();
        }
        // 1. 插入的节点设置为红色
        e.color = RED;
        // 2.对红黑树进行旋转和变色
        // 本质上就是父节点是黑色的就不需要调整，对应的 2 3的情况
        while (e != null && e != root && e.parent.color == RED) {
            // 找到所有的需要调整处理【旋转或者变色】的情况
            if (parentOf(e) == parentOf(parentOf(e)).left) {
                // 当前节点的父节点是爷爷节点的左子节点
                // 获取当前节点的叔叔节点，即爷爷节点的右节点
                RBNode y = rightOf(parentOf(parentOf(e)));
                if (colorOf(y) == RED) {
                    // 说明当前需要插入的节点的叔叔节点是存在的。那么就是上面的情况3
                    // 变色处理  父亲节点和叔叔节点变为黑色。爷爷节点变成红色
                    setColor(parentOf(e), BLACK);
                    setColor(y, BLACK);
                    setColor(parentOf(parentOf(e)), RED);
                    // 需要做一个递归处理
                    e = parentOf(parentOf(e));
                } else {
                    // 说明当前需要插入的节点的叔叔节点是不存在的。那么就是上面的情况2
                    if (e == parentOf(e).right) {
                        // 说明要做一次左旋的操作
                        e = parentOf(e);
                        leftRotate(e);
                    }
                    // 爷爷节点做右旋操作
                    setColor(parentOf(e), BLACK);
                    setColor(parentOf(parentOf(e)), RED);
                    rightRotate(parentOf(parentOf(e)));
                }
            } else {
                // 当前节点的父节点是爷爷节点的右子节点
                // 获取当前节点的叔叔节点，即爷爷节点的左节点
                RBNode x = leftOf(parentOf(parentOf(e)));
                if (colorOf(x) == RED) {
                    // 说明当前需要插入的节点的叔叔节点是存在的。那么就是上面的情况3
                    // 变色处理  父亲节点和叔叔节点变为黑色。爷爷节点变成红色
                    setColor(parentOf(e), BLACK);
                    setColor(x, BLACK);
                    setColor(parentOf(parentOf(e)), RED);
                    // 需要做一个递归处理
                    e = parentOf(parentOf(e));
                } else {
                    // 说明当前需要插入的节点的叔叔节点是不存在的。那么就是上面的情况2
                    if (e == parentOf(e).left) {
                        // 说明要做一次右旋的操作
                        e = parentOf(e);
                        rightRotate(e);
                    }
                    // 爷爷节点做左旋操作
                    setColor(parentOf(e), BLACK);
                    setColor(parentOf(parentOf(e)), RED);
                    leftRotate(parentOf(parentOf(e)));
                }
            }
        }
        // 3.根节点设置为黑色
        this.root.color = BLACK;
    }

    private boolean colorOf(RBNode node) {
        return node == null ? BLACK : node.color;
    }

    private RBNode parentOf(RBNode node) {
        return node != null ? node.parent : null;
    }

    private RBNode leftOf(RBNode node) {
        return node != null ? node.left : null;
    }

    private RBNode rightOf(RBNode node) {
        return node != null ? node.right : null;
    }

    private void setColor(RBNode node, boolean color) {
        if (node != null) {
            node.setColor(color);
        }
    }

    /**
     * 红黑树中的节点
     * 自定义节点
     * key 需要能比较
     */
    static class RBNode<K extends Comparable<K>, V> {
        // 当前节点的父节点
        private RBNode parent;
        // 当前节点的左节点
        private RBNode left;
        // 当前节点的右节点
        private RBNode right;
        // 当前节点的颜色(红F、黑T)
        private boolean color;
        private K key;

        private V value;

        public RBNode() {
        }

        public RBNode(K key, V value, boolean color) {
            this.color = color;
            this.key = key;
            this.value = value;
        }

        public RBNode(K key, V value, RBNode parent) {
            this.parent = parent;
            this.key = key;
            this.value = value;
        }


        public RBNode(RBNode parent, RBNode left, RBNode right, boolean color, K key, V value) {
            this.parent = parent;
            this.left = left;
            this.right = right;
            this.color = color;
            this.key = key;
            this.value = value;
        }

        public RBNode getParent() {
            return parent;
        }

        public void setParent(RBNode parent) {
            this.parent = parent;
        }

        public RBNode getLeft() {
            return left;
        }

        public void setLeft(RBNode left) {
            this.left = left;
        }

        public RBNode getRight() {
            return right;
        }

        public void setRight(RBNode right) {
            this.right = right;
        }

        public boolean isColor() {
            return color;
        }

        public void setColor(boolean color) {
            this.color = color;
        }

        public K getKey() {
            return key;
        }

        public void setKey(K key) {
            this.key = key;
        }

        public V getValue() {
            return value;
        }

        public void setValue(V value) {
            this.value = value;
        }
    }
}
